{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import scipy.misc\n",
    "import scipy.io\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.framework import ops\n",
    "ops.reset_default_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Start a graph session\n",
    "sess = tf.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Image Files\n",
    "original_image_file = 'temp/book_cover.jpg'\n",
    "style_image_file = 'temp/starry_night.jpg'\n",
    "\n",
    "# Saved VGG Network path\n",
    "vgg_path = '/home/nick/Documents/tensorflow/vgg_19_models/imagenet-vgg-verydeep-19.mat'\n",
    "\n",
    "# Default Arguments\n",
    "original_image_weight = 0.0005\n",
    "style_image_weight = 0.02\n",
    "regularization_weight = 50.0\n",
    "learning_rate = 0.01\n",
    "generations = 20\n",
    "output_generations = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Read in images\n",
    "original_image = scipy.misc.imread(original_image_file)\n",
    "style_image = scipy.misc.imread(style_image_file)\n",
    "\n",
    "# Get shape of target and make the style image the same\n",
    "target_shape = original_image.shape\n",
    "style_image = scipy.misc.imresize(style_image, float(target_shape[1]) / float(style_image.shape[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# VGG-19 Layer Setup\n",
    "# From paper\n",
    "vgg_layers = ['conv1_1', 'relu1_1',\n",
    "              'conv1_2', 'relu1_2', 'pool1',\n",
    "              'conv2_1', 'relu2_1',\n",
    "              'conv2_2', 'relu2_2', 'pool2',\n",
    "              'conv3_1', 'relu3_1',\n",
    "              'conv3_2', 'relu3_2',\n",
    "              'conv3_3', 'relu3_3',\n",
    "              'conv3_4', 'relu3_4', 'pool3',\n",
    "              'conv4_1', 'relu4_1',\n",
    "              'conv4_2', 'relu4_2',\n",
    "              'conv4_3', 'relu4_3',\n",
    "              'conv4_4', 'relu4_4', 'pool4',\n",
    "              'conv5_1', 'relu5_1',\n",
    "              'conv5_2', 'relu5_2',\n",
    "              'conv5_3', 'relu5_3',\n",
    "              'conv5_4', 'relu5_4']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Extract weights and matrix means\n",
    "def extract_net_info(path_to_params):\n",
    "    vgg_data = scipy.io.loadmat(path_to_params)\n",
    "    normalization_matrix = vgg_data['normalization'][0][0][0]\n",
    "    mat_mean = np.mean(normalization_matrix, axis=(0,1))\n",
    "    network_weights = vgg_data['layers'][0]\n",
    "    return(mat_mean, network_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Create the VGG-19 Network\n",
    "def vgg_network(network_weights, init_image):\n",
    "    network = {}\n",
    "    image = init_image\n",
    "\n",
    "    for i, layer in enumerate(vgg_layers):\n",
    "        if layer[1] == 'c':\n",
    "            weights, bias = network_weights[i][0][0][0][0]\n",
    "            weights = np.transpose(weights, (1, 0, 2, 3))\n",
    "            bias = bias.reshape(-1)\n",
    "            conv_layer = tf.nn.conv2d(image, tf.constant(weights), (1, 1, 1, 1), 'SAME')\n",
    "            image = tf.nn.bias_add(conv_layer, bias)\n",
    "        elif layer[1] == 'r':\n",
    "            image = tf.nn.relu(image)\n",
    "        else:\n",
    "            image = tf.nn.max_pool(image, (1, 2, 2, 1), (1, 2, 2, 1), 'SAME')\n",
    "        network[layer] = image\n",
    "    return(network)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Here we define which layers apply to the original or style image\n",
    "original_layer = 'relu4_2'\n",
    "style_layers = ['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1', 'relu5_1']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Get network parameters\n",
    "normalization_mean, network_weights = extract_net_info(vgg_path)\n",
    "\n",
    "shape = (1,) + original_image.shape\n",
    "style_shape = (1,) + style_image.shape\n",
    "original_features = {}\n",
    "style_features = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Get network parameters\n",
    "image = tf.placeholder('float', shape=shape)\n",
    "vgg_net = vgg_network(network_weights, image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Normalize original image\n",
    "original_minus_mean = original_image - normalization_mean\n",
    "original_norm = np.array([original_minus_mean])\n",
    "original_features[original_layer] = sess.run(vgg_net[original_layer], feed_dict={image: original_norm})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Get style image network\n",
    "image = tf.placeholder('float', shape=style_shape)\n",
    "vgg_net = vgg_network(network_weights, image)\n",
    "style_minus_mean = style_image - normalization_mean\n",
    "style_norm = np.array([style_minus_mean])\n",
    "\n",
    "for layer in style_layers:\n",
    "    layer_output = sess.run(vgg_net[layer], feed_dict={image: style_norm})\n",
    "    layer_output = np.reshape(layer_output, (-1, layer_output.shape[3]))\n",
    "    style_gram_matrix = np.matmul(layer_output.T, layer_output) / layer_output.size\n",
    "    style_features[layer] = style_gram_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Make Combined Image\n",
    "initial = tf.random_normal(shape) * 0.05\n",
    "image = tf.Variable(initial)\n",
    "vgg_net = vgg_network(network_weights, image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Loss\n",
    "original_loss = original_image_weight * (2 * tf.nn.l2_loss(vgg_net[original_layer] - original_features[original_layer]) /\n",
    "                original_features[original_layer].size)\n",
    "\n",
    "# Loss from Style Image\n",
    "style_loss = 0\n",
    "style_losses = []\n",
    "for style_layer in style_layers:\n",
    "    layer = vgg_net[style_layer]\n",
    "    feats, height, width, channels = [x.value for x in layer.get_shape()]\n",
    "    size = height * width * channels\n",
    "    features = tf.reshape(layer, (-1, channels))\n",
    "    style_gram_matrix = tf.matmul(tf.transpose(features), features) / size\n",
    "    style_expected = style_features[style_layer]\n",
    "    style_losses.append(2 * tf.nn.l2_loss(style_gram_matrix - style_expected) / style_expected.size)\n",
    "style_loss += style_image_weight * tf.reduce_sum(style_losses)\n",
    "\n",
    "# To Smooth the resuts, we add in total variation loss       \n",
    "total_var_x = sess.run(tf.reduce_prod(image[:,1:,:,:].get_shape()))\n",
    "total_var_y = sess.run(tf.reduce_prod(image[:,:,1:,:].get_shape()))\n",
    "first_term = regularization_weight * 2\n",
    "second_term_numerator = tf.nn.l2_loss(image[:,1:,:,:] - image[:,:shape[1]-1,:,:])\n",
    "second_term = second_term_numerator / total_var_y\n",
    "third_term = (tf.nn.l2_loss(image[:,:,1:,:] - image[:,:,:shape[2]-1,:]) / total_var_x)\n",
    "total_variation_loss = first_term * (second_term + third_term)\n",
    "\n",
    "# Combined Loss\n",
    "loss = original_loss + style_loss + total_variation_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Declare Optimization Algorithm\n",
    "optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "train_step = optimizer.minimize(loss)\n",
    "\n",
    "# Initialize Variables and start Training\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generation 1 out of 20\n",
      "Generation 2 out of 20\n",
      "Generation 3 out of 20\n",
      "Generation 4 out of 20\n",
      "Generation 5 out of 20\n",
      "Generation 6 out of 20\n",
      "Generation 7 out of 20\n",
      "Generation 8 out of 20\n",
      "Generation 9 out of 20\n",
      "Generation 10 out of 20\n",
      "Generation 11 out of 20\n",
      "Generation 12 out of 20\n",
      "Generation 13 out of 20\n",
      "Generation 14 out of 20\n",
      "Generation 15 out of 20\n",
      "Generation 16 out of 20\n",
      "Generation 17 out of 20\n",
      "Generation 18 out of 20\n",
      "Generation 19 out of 20\n",
      "Generation 20 out of 20\n"
     ]
    }
   ],
   "source": [
    "# Declare Optimization Algorithm\n",
    "optimizer = tf.train.GradientDescentOptimizer(learning_rate)\n",
    "train_step = optimizer.minimize(loss)\n",
    "\n",
    "# Initialize Variables and start Training\n",
    "sess.run(tf.global_variables_initializer())\n",
    "for i in range(generations):\n",
    "    \n",
    "    sess.run(train_step)\n",
    "\n",
    "    # Print update and save temporary output\n",
    "    if (i+1) % output_generations == 0:\n",
    "        print('Generation {} out of {}'.format(i + 1, generations))\n",
    "        image_eval = sess.run(image)\n",
    "        best_image_add_mean = image_eval.reshape(shape[1:]) + normalization_mean\n",
    "        output_file = 'temp_output_{}.jpg'.format(i)\n",
    "        scipy.misc.imsave(output_file, best_image_add_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
